# -*- coding:utf-8 -*-

import pymysql
import configparser

class MySQLDatabase:
    def __init__(self, config_file='/home/ubuntu/code/crawl-annual-report/conf/db_config.ini'):
        """
        初始化数据库连接参数，并读取配置文件。
        :param config_file: 配置文件路径，默认为 'db_config.ini'
        """
        self.config_file = config_file
        self.host = None
        self.user = None
        self.password = None
        self.database = None
        self.charset = None
        self.connection = None
        self.cursor = None
        self.read_config()

    def read_config(self):
        """
        读取配置文件中的数据库连接信息。
        """
        config = configparser.ConfigParser()
        config.read(self.config_file, encoding='utf-8')
        if 'mysql' in config:
            self.host = config.get('mysql', 'host')
            self.user = config.get('mysql', 'user')
            self.password = config.get('mysql', 'password')
            self.database = config.get('mysql', 'database')
            self.charset = config.get('mysql', 'charset', fallback='utf8mb4')
        else:
            raise ValueError(f"配置文件中未找到 'mysql' 部分。")

    def connect(self):
        """
        建立数据库连接。
        """
        if self.connection is None:
            self.connection = pymysql.connect(
                host=self.host,
                user=self.user,
                password=self.password,
                database=self.database,
                charset=self.charset
            )
            self.cursor = self.connection.cursor()
            print("数据库连接成功。")
        else:
            print("数据库已连接。")

    def close(self):
        """
        关闭数据库连接。
        """
        if self.cursor:
            self.cursor.close()
        if self.connection:
            self.connection.close()
        print("数据库连接已关闭。")

    def create_table(self):
        """
        创建数据表（如果尚未存在）。
        """
        create_table_sql = """
        CREATE TABLE IF NOT EXISTS listed_company_report (
            id INT AUTO_INCREMENT PRIMARY KEY,
            company_code VARCHAR(50) NOT NULL,
            company_name VARCHAR(255) NOT NULL,
            publish_date DATE DEFAULT NULL,
            work_report TEXT NOT NULL
        ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;
        """
        self.cursor.execute(create_table_sql)
        self.connection.commit()
        print("数据表创建成功。")

    def insert_data(self, data):
        """
        批量插入数据，如果数据已存在则跳过插入。
        :param data: 包含要插入数据的列表，每个元素是一个元组，对应一行数据
        """
        insert_sql = """
        INSERT INTO listed_company (id, company_code, company_name, org_id, announcement_id, announcement_time, announcement_title, base_url, subset_url, download_success)
        VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
        """

        # 检查是否已存在的 SQL 语句
        check_sql = """
        SELECT COUNT(*) FROM listed_company
        WHERE company_code = %s AND announcement_time = %s and announcement_title = %s;
        """
        
        for record in data:
            company_code = record[1]
            announcement_time = record[4]
            announcement_title = record[5]
            
            # 查询是否已经存在该记录
            self.cursor.execute(check_sql, (company_code, announcement_time, announcement_title))
            result = self.cursor.fetchone()

            # 如果查询结果是 0，表示没有该记录
            if result[0] == 0:
                self.cursor.execute(insert_sql, record)
                print(f"插入数据: {record}")
            else:
                print(f"数据已存在，跳过插入: {record}")
        
        self.connection.commit()
        print("数据插入完成。")

    def _insert_data(self, insert_sql, data):
        """
        批量插入数据，如果数据已存在则跳过插入。
        :param data: 包含要插入数据的列表，每个元素是一个元组，对应一行数据
        """

        # 检查是否已存在的 SQL 语句
        check_sql = """
        SELECT COUNT(*) FROM listed_company_report
        WHERE id = %s;
        """
        
        for record in data:
            id = record[0]
            
            # 查询是否已经存在该记录
            self.cursor.execute(check_sql, (id,))
            result = self.cursor.fetchone()

            # 如果查询结果是 0，表示没有该记录
            if result[0] == 0:
                self.cursor.execute(insert_sql, record)
                print(f"插入数据: {record}")
            else:
                print(f"有该记录，不用插入: {record}")
        
        self.connection.commit()
        print("数据插入完成。")

    def query_data(self, sql):
        """
        查询数据
        :param sql: 查询语句
        :param params: 查询参数
        :return: 查询结果
        """
        try:
            # self.connect()
            with self.connection.cursor() as cursor:
                cursor.execute(sql)
                result = cursor.fetchall()
                columns = [desc[0] for desc in cursor.description]
                return result, columns
        except Exception as e:
            print(f"查询数据时出错: {e}")
            return None
        # finally:
        #     self.close()

    def update_data(self, sql):
        """
        更新数据
        :param sql: 更新语句
        :param params: 更新参数
        :return: 受影响的行数
        """
        try:
            # self.connect()
            with self.connection.cursor() as cursor:
                affected_rows = cursor.execute(sql)
                self.connection.commit()
                return affected_rows
        except Exception as e:
            print(f"更新数据时出错: {e}")
            return 0
        # finally:
        #     self.close()

# 使用示例
if __name__ == '__main__':
    # 创建数据库对象
    db = MySQLDatabase(config_file='db_config.ini')
    # 建立连接
    db.connect()
    # 创建数据表
    db.create_table()
    # 准备要插入的数据
    data_to_insert = [
        ('000001', '公司A', '2025-01-01', '报告内容A'),
        ('000002', '公司B', '2025-02-01', '报告内容B'),
        ('000003', '公司C', None, '报告内容C'),
        ('000001', '公司A', '2025-01-01', '报告内容A')  # 重复的数据，应该跳过
    ]
    # 插入数据
    db.insert_data(data_to_insert)
    # 关闭连接
    db.close()
